#include "kernel_operator.h"
#include <type_traits>
#include <cmath>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename T> class KernelElements {
public:
    __aicore__ inline KernelElements() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, 
                            int32_t var_dim[3], int32_t indices_dim[3], int32_t updates_dim[3],
                            int32_t reduce) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->var_dim[0] = var_dim[0];
        this->var_dim[1] = var_dim[1];
        this->var_dim[2] = var_dim[2];
        this->indices_dim[0] = indices_dim[0];
        this->indices_dim[1] = indices_dim[1];
        this->indices_dim[2] = indices_dim[2];
        this->updates_dim[0] = updates_dim[0];
        this->updates_dim[1] = updates_dim[1];
        this->updates_dim[2] = updates_dim[2];
        this->reduce = reduce;

        varGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_VAR *>(var), 1);
        indicesGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_INDICES *>(indices), 1);
        updatesGm.SetGlobalBuffer(reinterpret_cast<__gm__ DTYPE_UPDATES *>(updates), 1);

        pipe.InitBuffer(QueueTmp1, 32 * sizeof(DTYPE_VAR));
        pipe.InitBuffer(QueueTmp2, 32 * sizeof(float));
    }
    //this->interval = 1  00 00 01 01 02 02 10 10 11 11 12 14
    __aicore__ inline void Process() {
        if(this->reduce == 1) //add Case1 Case4
        {
            for (int32_t i = 0; i < this->indices_dim[0]; i ++)
            {
                for (int32_t j = 0; j < this->indices_dim[1]; j ++)
                {
                    for (int32_t k = 0; k < this->indices_dim[2]; k ++)
                    {
                        DTYPE_INDICES indice = indicesGm.GetValue(i*this->indices_dim[1]*this->indices_dim[2] + j*this->indices_dim[2] + k);
                        DTYPE_VAR var = varGm.GetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k);
                        DTYPE_UPDATES update = updatesGm.GetValue(i*this->updates_dim[1]*this->updates_dim[2] + j*this->updates_dim[2] + k);
                        if constexpr (std::is_same_v<T, half>) {
                            auto tmp1 = QueueTmp1.Get<DTYPE_VAR>();
                            auto tmp2 = QueueTmp2.Get<DTYPE_VAR>();
                            tmp1.SetValue(0, var);
                            tmp2.SetValue(0, update);
                            Add(tmp1,tmp1,tmp2,1);
                            DTYPE_VAR resHalf = tmp1.GetValue(0);

                            // float varF32,updateF32;
                            // auto tmp1 = QueueTmp1.Get<DTYPE_VAR>();
                            // auto tmp2 = QueueTmp2.Get<float>();
                            // tmp1.SetValue(0, var);
                            // Cast(tmp2, tmp1, RoundMode::CAST_NONE, 1);
                            // varF32 = tmp2.GetValue(0);
                            // tmp1.SetValue(0, update);
                            // Cast(tmp2, tmp1, RoundMode::CAST_NONE, 1);
                            // updateF32 = tmp2.GetValue(0);
                            // half resHalf = (half)(varF32 + updateF32);

                            // // half转float
                            // int16_t varInt16 = *(int16_t*)&var;
                            // int32_t fltInt32 = ((varInt16 & 0x8000) << 16);
                            // fltInt32 |= ((varInt16 & 0x7fff) << 13) + 0x38000000;
                            // float varFp32 = *(float*)&fltInt32;

                            // int16_t updateInt16 = *(int16_t*)&update;
                            // fltInt32 = ((updateInt16 & 0x8000) << 16);
                            // fltInt32 |= ((updateInt16 & 0x7fff) << 13) + 0x38000000;
                            // float updateFp32 = *(float*)&fltInt32;
                            // float res = varFp32 + updateFp32;

                            // // float转half
                            // int16_t fltInt16;
                            // fltInt32 = *(int32_t*)&res;
                            // fltInt16 = ((fltInt32 & 0x7fffffff) >> 13) - (0x38000000 >> 13);
                            // fltInt16 |= ((fltInt32 & 0x80000000) >> 16);
                            // half resHalf = *(half*)&fltInt16;
                            varGm.SetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k, resHalf);
                        }
                        else
                        {
                            T resHalf = (T)var + (T)update;
                            varGm.SetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k, resHalf);
                        }
                    }
                }
            }
        }
        else if(this->reduce == 0) //none Case2 Case5
        {
            for (int32_t i = 0; i < this->indices_dim[0]; i ++)
            {
                for (int32_t j = 0; j < this->indices_dim[1]; j ++)
                {
                    for (int32_t k = 0; k < this->indices_dim[2]; k ++)
                    {
                        DTYPE_INDICES indice = indicesGm.GetValue(i*this->indices_dim[1]*this->indices_dim[2] + j*this->indices_dim[2] + k);
                        DTYPE_UPDATES update = updatesGm.GetValue(i*this->updates_dim[1]*this->updates_dim[2] + j*this->updates_dim[2] + k);
                        varGm.SetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k, update);
                    }
                }
            }//8 021 1 011 4 
            // varGm.SetValue(0, (DTYPE_VAR)this->cycles);
            // varGm.SetValue(1, (DTYPE_VAR)this->interval);
            // varGm.SetValue(2, (DTYPE_VAR)this->loopCount);
        }
        else if(this->reduce == 2) //multiply Case3
        {
            for (int32_t i = 0; i < this->indices_dim[0]; i ++)
            {
                for (int32_t j = 0; j < this->indices_dim[1]; j ++)
                {
                    for (int32_t k = 0; k < this->indices_dim[2]; k ++)
                    {
                        DTYPE_INDICES indice = indicesGm.GetValue(i*this->indices_dim[1]*this->indices_dim[2] + j*this->indices_dim[2] + k);
                        DTYPE_VAR var = varGm.GetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k);
                        DTYPE_UPDATES update = updatesGm.GetValue(i*this->updates_dim[1]*this->updates_dim[2] + j*this->updates_dim[2] + k);
                        if constexpr (std::is_same_v<T, half>) {
                            auto tmp1 = QueueTmp1.Get<DTYPE_VAR>();
                            auto tmp2 = QueueTmp2.Get<DTYPE_VAR>();
                            tmp1.SetValue(0, var);
                            tmp2.SetValue(0, update);
                            Mul(tmp1,tmp1,tmp2,1);
                            DTYPE_VAR resHalf = tmp1.GetValue(0);
                            varGm.SetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k, resHalf);
                        }
                        else
                        {
                            T resHalf = (T)var * (T)update;
                            varGm.SetValue(i*this->var_dim[1]*this->var_dim[2] + indice*this->var_dim[2] + k, resHalf);
                        }
                    }
                }
            }
        }
        // else
        //add Case1 Case4
        //none Case2 Case5
        //multiply Case3
        //half Case1
        //float Case2 Case5
        //int32 Case3

        //case1 add  half
        //case2 none float
        //case3 mult int32
        //case4 add  uint8
        //case5 none float

        // float temp_sum;
        // float temp_add;
        // auto tmp1 = QueueTmp1.Get<float>();
        if constexpr (std::is_same_v<T, int32_t>) 
        {
            TBuf<QuePosition::VECCALC> boox;                
            pipe.InitBuffer(boox, 9999999 * sizeof(float));
            auto booxxxxx = boox.Get<float>();
            Add(booxxxxx, booxxxxx, booxxxxx, 9999999);
            Add(booxxxxx[1], booxxxxx[2], booxxxxx[3], 9999999);
        }
        // for(int z = 0; z < this->loopCount; z++)
        // {
        //     int32_t x_num = z/this->interval;
        //     x_num = x_num*this->cycles*this->interval + z%this->interval;
        //     for(int i = 0; i < this->cycles; i++)
        //     {
        //         int32_t temp_num = x_num + i*this->interval;
        //         temp_add = xGm.GetValue(temp_num);
        //         tmp1.SetValue(0, (float)temp_add);
        //         Exp(tmp1, tmp1, 1);
        //         temp_add = tmp1.GetValue(0);

        //         if(i == 0)
        //         {
        //             temp_sum = (float)temp_add;
        //         }
        //         else
        //         {
        //             temp_sum = (float)temp_sum + (float)temp_add;
        //         }                
        //     }

        //     for(int i = 0; i < this->cycles; i++)
        //     {
        //         int32_t temp_num = x_num + i*this->interval;
        //         temp_add = xGm.GetValue(temp_num);
        //         tmp1.SetValue(0, (float)temp_add);
        //         Exp(tmp1, tmp1, 1);
        //         temp_add = tmp1.GetValue(0);
        //         temp_add = temp_add / temp_sum;
        //         yGm.SetValue(temp_num, (DTYPE_Y)temp_add);
        //     }
        // }
        // varGm.SetValue(0, (DTYPE_VAR)this->cycles);
        // // yGm.SetValue(1, (DTYPE_Y)this->interval);
        // // yGm.SetValue(2, (DTYPE_Y)this->loopCount);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIndices, inQueueUpdates;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueVar;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2;
    // TBuf<QuePosition::VECCALC> QueueBuff;
    
    GlobalTensor<DTYPE_VAR> varGm;
    GlobalTensor<DTYPE_INDICES> indicesGm;
    GlobalTensor<DTYPE_UPDATES> updatesGm;

    int32_t var_dim[3];
    int32_t indices_dim[3];
    int32_t updates_dim[3];

    int32_t reduce;
};
extern "C" __global__ __aicore__ void scatter_elements(GM_ADDR var, GM_ADDR indices, GM_ADDR updates, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelElements<DTYPE_VAR> op;
    op.Init(var, indices, updates,
                tiling_data.var_dim, tiling_data.indices_dim, tiling_data.updates_dim,
                tiling_data.reduce);
    op.Process();
}